import cv2
import numpy as np
import json
import logging
import os.path

logger = logging.getLogger('cardScanner')

logger.setLevel(logging.DEBUG)
logfile = logging.FileHandler(os.path.join(os.getcwd(), 'log.txt'), 'w')
logger.addHandler(logfile)


def read_template(tfile):
    with open(tfile, 'r') as infile:
        template = json.load(infile)
    return template

def read_abmark(file):
    with open(file, 'r') as infile:
        abmark = json.load(infile)
    return abmark

def abcheck(img, abmark):
    size = abmark['size']
    w = size['w']
    h = size['h']
    if h != img.shape[0] or w != img.shape[1]:
        img = cv2.resize(img, (w, h))
    mark = abmark['mark']
    x = abmark['x']
    y = abmark['y']
    cardtype = abmark['type']
    assert(cardtype in "AB")
    if cardtype == "A":
        nonetype = "B"
    else:
        nonetype = "A"
    center_x = x + mark['w']//2
    center_y = y + mark['h']//2
    """
    cv2.circle(img, (center_x, center_y), 3, color=255)
    cv2.imshow('col', img)
    cv2.waitKey(-1)
    cv2.destroyAllWindows()
    """
    return cardtype if check_mark(img, center_x, center_y, r=3, tol=120) else nonetype

def check_mark(img, x, y, r = 2, tol = 100):
    # select a few points decide if its marked
    area = img[y-r:y+r, x-r:x+r]
    area = area.reshape(-1)
    avg = np.mean(area)
    return avg <= tol


def split_columns(img, n):
    width = img.shape[1]
    col_width = width // n
    list_of_columns = []
    for i in range(n):
        start = i * col_width
        list_of_columns.append(img[:, start:start + col_width])
    return list_of_columns


def to_center(pos, w, h):
    # get the center position
    return pos[0] + w // 2, pos[1] + h // 2


def find_block(column_image, mark):
    w = mark['w']
    h = mark['h']
    block = np.ones([h, w], dtype=np.uint8)
    res = cv2.matchTemplate(column_image, block, cv2.TM_SQDIFF)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
    return to_center(min_loc, w, h)


def read_column(column_image, answer_map, mark):
    x, y = find_block(column_image, mark)
    img = column_image.copy()
    if logger.level >= logging.INFO:
        cv2.circle(img, (x, y), 3, color=255)
        cv2.imshow('col', img)
        cv2.waitKey(1000)
        cv2.destroyAllWindows()
    h = mark['h'] * 1.6
    for result in answer_map:
        y_min = answer_map[result]
        if y <= y_min + h and y >= y_min:
            return result
    else:
        # if failed use method b
        return read_column_b(column_image, answer_map, mark)

def read_column_b(image, answer_map, mark):
    h, w = image.shape
    mark_h = mark['h']
    center_x = w//2
    logger.debug('plan b')
    for result in answer_map:
        y = answer_map[result]
        center_y = int(y + mark_h//2*1.6)
        if check_mark(image, center_x, center_y, r= 3):
            debug_circel(image, center_x, center_y)
            return result
    else:
        return None


def debug_circel(image, x, y, enable = False):
    if enable:
        image = image.copy()
        cv2.circle(image, (x, y), 2, 255, -1)
        cv2.imshow('debug', image)
        cv2.waitKey(-1)
        cv2.destroyAllWindows()

def get_block(img, top_left, block_info):
    x, y = top_left
    w = block_info['w']
    h = block_info['h']
    return img[y:y + h, x:x + w]


def read_sheet(img_gray, template_info):
    size = template_info['size']
    w = size['w']
    h = size['h']
    if h != img_gray.shape[0] or w != img_gray.shape[1]:
        img_gray = cv2.resize(img_gray, (w, h))
    blocks = template_info['blocks']
    block_info = template_info['block_info']
    mark = template_info['mark']

    answer_map = template_info['answer_map']

    answers = []

    for block in blocks:
        block_img = get_block(img_gray, block, block_info)
        for col_img in split_columns(block_img, block_info['n']):
            answers.append(read_column(col_img, answer_map, mark))
    if 'student_num' in template_info:
        student_num = template_info['student_num']
        snum = read_student_num(img_gray, student_num, mark)
    else:
        snum = ''
    return snum, answers


def read_student_num(img, student_num_info, mark):
    x = student_num_info['x']
    y = student_num_info['y']
    block_img = get_block(img, (x,y),student_num_info)
    nums = []
    num_map = student_num_info['num_map']
    for col_img in split_columns(block_img, student_num_info['n']):
        nums.append(read_column(col_img, num_map, mark))
    return convert_compact_form(nums)

def convert_compact_form(answers):
    return "".join(c if c else '#' for c in answers)


if __name__ == '__main__':
    template_info = read_template('cards/eighty-twk2.json')
    answer_sheet = 'cards/test.jpg'
    #answer_sheet = 'cards/eighty.png'
    sheet_img = cv2.imread(answer_sheet, cv2.IMREAD_COLOR)
    sheet_img_gray = cv2.cvtColor(sheet_img, cv2.COLOR_RGB2GRAY)
    student_num, answers = read_sheet(sheet_img_gray, template_info)
    print(student_num,convert_compact_form(answers))
